#include "solver.h"

Solver::Solver() {
	_Y0=_I0=_Q0=_Y1=_Y2=_I1=_I2=_Q1=_Q2=NULL;
	m_row=m_col=0; m_timeBegin=m_timeEnd=0;
}

Solver::Solver ( Mat &p_inputImage, Mat &p_mask ){
	_Y0=_I0=_Q0=_Y1=_Y2=_I1=_I2=_Q1=_Q2=NULL;
	m_inputImage=p_inputImage.clone();
	m_mask=p_mask.clone();
	m_outputImage=p_mask.clone();
	m_row=m_inputImage.size().height;
	m_col=m_inputImage.size().width;
	m_timeBegin=m_timeEnd=0;
}

Solver::~Solver () {
	if(_Y0) delete _Y0; _Y0=NULL;
	if(_Y1) delete _Y1; _Y1=NULL;
	if(_Y2) delete _Y2; _Y2=NULL;
	if(_I0) delete _I0; _I0=NULL;
	if(_I1) delete _I1; _I1=NULL;
	if(_I2) delete _I2; _I2=NULL;
	if(_Q0) delete _Q0; _Q0=NULL;
	if(_Q1) delete _Q1; _Q1=NULL;
	if(_Q2) delete _Q2; _Q2=NULL;
}

int Solver::setInputImage(Mat &p) {
	m_inputImage=p.clone();
	resize();
	return 1;
}

int Solver::setMaskImage(Mat &p) {
	m_mask=p.clone();
	return 1;
}

Mat &Solver::getOutputImage() {
	return m_outputImage;
}

int Solver::resize() {
	m_row=m_inputImage.size().height;
	m_col=m_inputImage.size().width;
	return 1;
}

double Solver::getSolvingTime() {
	return (m_timeEnd-m_timeBegin)/1000.0;//difftime(m_timeEnd,m_timeBegin)/1000.0;
}

double Solver::normalPdf ( double p_mean, double p_variance, double p_x ) {
	return exp(-(p_x*p_x)/(p_variance*2));
}

int Solver::RGBnormalize ( double &r, double &g, double &b ) {

	double t_a, t_b;

	t_a=min(r,min(g,b));
	t_b=max(r,max(g,b));

	if(t_a<0&&t_b>255)
	{
		r=(r-t_a)/(t_b-t_a)*255;
		g=(g-t_a)/(t_b-t_a)*255;
		b=(b-t_a)/(t_b-t_a)*255;
	}
	else if(t_a<0&&t_b<255+EPS)
	{
		r=(r-t_a)/(t_b-t_a)*t_b;
		g=(g-t_a)/(t_b-t_a)*t_b;
		b=(b-t_a)/(t_b-t_a)*t_b;
	}
	else if(t_a>0-EPS&&t_b>255)
	{
		r=(r-t_a)/(t_b-t_a)*(255-t_a)+t_a;
		g=(g-t_a)/(t_b-t_a)*(255-t_a)+t_a;
		b=(b-t_a)/(t_b-t_a)*(255-t_a)+t_a;
	}

	return 1;
}

int Solver::mat2Array ( Mat &m, double **p_y, double **p_i, double **p_q ) {

	int i,j,k,cnt,R,G,B,r,c;
	double Y,I,Q,*t_y,*t_i,*t_q;

	r=getRow(); c=getCol();
	
	t_y = new double[r*c]; memset(t_y,0,sizeof(t_y));
	t_i = new double[r*c]; memset(t_i,0,sizeof(t_i));
	t_q = new double[r*c]; memset(t_q,0,sizeof(t_q));

	unsigned char *p = m.data;

	for(i=cnt=k=0;i<m_row;i++)
		for(j=0;j<m_col;j++)
		{
			B=p[cnt];
			G=p[cnt+1];
			R=p[cnt+2];
			Y=0.299*R+0.587*G+0.114*B;
			I=0.595716*R-0.274453*G-0.321263*B;
			Q=0.211456*R-0.522591*G+0.311135*B;
			t_y[k]=Y;
			t_i[k]=I;
			t_q[k]=Q;
			cnt+=3; k++;
		}

	(*p_y) = t_y;
	(*p_i) = t_i;
	(*p_q) = t_q;

	return 1;

}

int Solver::constructSparseMatrix () {

	int i,j,k,cnt,s,p,q,u,v,len,r,c,n;
	double t_val,x_1,x_2,mean,variance;
	r=getRow();
	c=getCol();
	n=r*c;

	int *pidx_r,*pidx_c,*idx_r,*idx_c;
	double *pval,*gval;
	int wd=1; int wdx=2*wd+1; wdx*=wdx;
	A=SparseMatrix(n,n,n*wdx);

	idx_r = new int[wdx];
	idx_c = new int[wdx];
	gval = new double[wdx+1];
	int mxLen=0;

	int *tidx_c = new int[n*wdx];
	int *tidx_len = new int[n];
	double *tidx_v = new double [n*wdx];
	memset(tidx_len,0,sizeof(int)*n);

	int tot=0;

	for(i=cnt=0;i<m_row;i++)
		for(j=0;j<m_col;j++)
		{
			s=abs(_Y1[cnt]-_Y2[cnt])+abs(_I1[cnt]-_I2[cnt])+abs(_Q1[cnt]-_Q2[cnt]);

			len=0; x_1=x_2=0;
			if(s<=3)
				for(p=-wd;p<=wd;p++)
					for(q=-wd;q<=wd;q++)
						if(p!=0||q!=0)
						{
							u=(i+p); v=(j+q);
							if(0<=u&&u<r&&0<=v&&v<c)
							{
								idx_r[len]=cnt;
								idx_c[len]=u*c+v;
								gval[len]=_Y1[u*c+v];
								x_1+=gval[len];
								x_2+=gval[len]*gval[len];
								len++;
							}
						}
			t_val=_Y1[cnt];
			gval[len+1]=t_val;
			x_1+=t_val;
			x_2+=t_val*t_val;

			mean=x_1/(1+len);
			variance=x_2/(1+len)-mean*mean;
			if(variance<1e-9)
				variance=1e-9;

			x_1=0;
			for(k=0;k<len;k++)
			{
				gval[k]=-normalPdf(0,variance,gval[k]-t_val);
				x_1-=gval[k];
			}
			for(k=0;k<len;k++)
				gval[k]/=x_1;
			idx_r[len]=idx_c[len]=cnt;
			gval[len]=1; len++;


			// indices are 0-based
			for(k=0;k<len;k++,tot++)
			{
				int &l=tidx_len[idx_c[k]];
				int &p_c=tidx_c[wdx*idx_c[k]+l];
				double &p_v=tidx_v[wdx*idx_c[k]+l];
				p_c=idx_r[k]; p_v=gval[k]; l++;
			}

			++cnt;
		}

	pidx_r = A.rep->r;
	pidx_c = A.rep->c;
	pval = A.rep->d;
	for(i=0,j=0;i<n;i++)
	{
		pidx_c[i]=j;
		for(k=0;k<tidx_len[i];k++,j++)
		{
			pidx_r[j]=tidx_c[wdx*i+k];
			pval[j]=tidx_v[wdx*i+k];
		}
	}
	pidx_c[n]=tot;
	// csr : 17s 35s
	// coo : 14s 39s
	// csc : 5s

	// 1024 x 768 = 10s * 2
	delete tidx_c; delete tidx_v; delete tidx_len;
	delete gval; delete idx_r; delete idx_c;

	return 1;

}

int Solver::constructColumnVectors () {
	int r,c;
	r=getRow(); c=getCol();

	b1=ColumnVector(r*c,0.0);
	b2=ColumnVector(r*c,0.0);

	int i,j,k;
	for(i=k=0;i<r;i++)
		for(j=0;j<c;j++,k++)
		{
			double a,b,c;
			a=_Y1[k]-_Y2[k];
			b=_I1[k]-_I2[k];
			c=_Q1[k]-_Q2[k];
			if(abs(a)+abs(b)+abs(c)>3)
			{
				b=_I2[k]; c=_Q2[k];
				b1(k)=b;
				b2(k)=c;
			}
		}

	return 1;

}

Mat& Solver::solve() {

	if(m_inputImage.size().area()>2e6)
		return solveReduced(1);

	mat2Array(m_inputImage,&_Y1,&_I1,&_Q1);
	mat2Array(m_mask,&_Y2,&_I2,&_Q2);
	constructSparseMatrix();
	constructColumnVectors();
	x1=A.solve(b1);
	x2=A.solve(b2);

	unsigned char *p = m_outputImage.data;

	int i,j,cnt,k,r,c;
	r=getRow(); c=getCol();

	for(i=cnt=k=0;i<r;i++)
		for(j=0;j<c;j++)
		{
			double R,G,B,Y,I,Q;
			Y=(int)_Y1[k];
			I=(int)x1(k);// if(I>153) I=153; if(I<-153) I=-153;
			Q=(int)x2(k);// if(Q>151) Q=151; if(Q<-151) Q=-151;
//			printf("%d %d %lf %lf %lf\n",i,j,Y,I,Q);
			R=1.0*Y+0.9563*I+0.6210*Q;
			G=1.0*Y-0.2721*I-0.6474*Q;
			B=1.0*Y-1.1070*I+1.7046*Q;
			RGBnormalize(R,G,B);
			p[cnt]=(int)B;
			p[cnt+1]=(int)G;
			p[cnt+2]=(int)R;
			cnt+=3; k++;
		}

	m_timeEnd=clock();//time(NULL)*1000;
	return m_outputImage;
}

Mat &Solver::solveReduced(int p_isReducedSize) {

	Mat s0; cv::resize(m_inputImage,s0,Size(0,0),0.5,0.5,INTER_NEAREST);
	Mat s1; cv::resize(m_mask,s1,Size(0,0),0.5,0.5,INTER_NEAREST);

	Solver *t_solver = new Solver();
	t_solver->solve(s0,s1,p_isReducedSize-1);

	Mat s2; cv::resize(t_solver->getOutputImage(),s2,m_inputImage.size(),0,0,INTER_LINEAR);

	unsigned char *p0 = m_inputImage.data;
	unsigned char *p1 = s2.data;
	unsigned char *p2 = m_outputImage.data;

	int i,j,cnt,r,c;
	r=getRow(); c=getCol();

	for(i=cnt=0;i<r;i++)
		for(j=0;j<c;j++)
		{
			double R,G,B,Y,I,Q;

			B=(int)p1[cnt];
			G=(int)p1[cnt+1];
			R=(int)p1[cnt+2];

			I=0.595716*R-0.274453*G-0.321263*B;
			Q=0.211456*R-0.522591*G+0.311135*B;

			B=p0[cnt];
			G=p0[cnt+1];
			R=p0[cnt+2];

			Y=0.299*R+0.587*G+0.114*B;
//			printf("%d %d %lf %lf %lf\n",i,j,Y,I,Q);
			R=1.0*Y+0.9563*I+0.6210*Q;
			G=1.0*Y-0.2721*I-0.6474*Q;
			B=1.0*Y-1.1070*I+1.7046*Q;

			RGBnormalize(R,G,B);
			p2[cnt]=(int)B;
			p2[cnt+1]=(int)G;
			p2[cnt+2]=(int)R;

			cnt+=3;
		}

	delete t_solver;

	m_timeEnd=clock();//time(NULL)*1000;
	return m_outputImage;

}

Mat &Solver::solve(Mat &p_inputImage, Mat &p_mask, int p_isReducedSize=0) {

	setInputImage(p_inputImage);
	setMaskImage(p_mask);
	m_outputImage=Mat(p_inputImage.size(),p_inputImage.type());
	resize();

	m_timeBegin=clock();//time(NULL)*1000;

	if(p_isReducedSize)
		return solveReduced(p_isReducedSize);
	else return solve();
}
